#include "websocket.h"

#include <QStringList>
#include <QRegExp>
#include <QCryptographicHash>

WebSocket::WebSocket(QObject *parent) : Socket(parent)
{
    v_state = Handshake;
}

qint64 WebSocket::write(const QByteArray& data)
{
    QByteArray toSend;
    int dataCount = data.count();
    toSend.append(129);
    if(dataCount <= 125)
       toSend.append((char)data.count());
    else if(dataCount <= 0xFFFF)
    {
        toSend.append(126);
        toSend.append((dataCount & 0xFF00) >> 8);
        toSend.append(dataCount & 0xFF);
    }else{
        toSend.append(127);
        toSend.append((dataCount & 0xFF00000000000000) >> 56);
        toSend.append((dataCount & 0xFF000000000000) >> 48);
        toSend.append((dataCount & 0xFF0000000000) >> 40);
        toSend.append((dataCount & 0xFF00000000) >> 32);
        toSend.append((dataCount & 0xFF000000) >> 24);
        toSend.append((dataCount & 0xFF0000) >> 16);
        toSend.append((dataCount & 0xFF00) >> 8);
        toSend.append(dataCount & 0xFF);
    }
    toSend += data;
    return QIODevice::write(toSend);
}


void WebSocket::processIncomingData()
{
    v_dataStream += this->readAll();
    QStringList requests;
    QString handshake;
    QRegExp keyRegexp("Sec-WebSocket-Key: (.+)\r\n");
    QRegExp protocolRegexp("Sec-WebSocket-Protocol: (.+)\r\n");
    keyRegexp.setMinimal(true);
    QString key;
    QString acceptKey;
    QByteArray decodedData;
    QByteArray encodedData;
    int maskCurrentIndex;
    Frame frame;

    switch(v_state){
    case Handshake:
        requests = QString(v_dataStream).split(QRegExp("\r\n\r\n"), QString::KeepEmptyParts);

	// counts the empty request
        if(requests.count() >= 2){
            v_dataStream = requests.at(1).toLatin1();
            handshake = requests.at(0);
            protocolRegexp.indexIn(handshake);
            qDebug() << "Protocole=" << protocolRegexp.cap(1);
            keyRegexp.indexIn(handshake);
            // Gets and stores the key
            key = keyRegexp.capturedTexts().at(1);
            if(key.isEmpty()){
                Q_EMIT handshakeFailed();
                this->close();
                break;
            }
            acceptKey = computeAcceptKey(key);
        }else{
            this->close();
            break;
        }

        // Replies to the request to accept the connection
        QIODevice::write(QString("HTTP/1.1 101 Switching Protocols\r\n"
                                "Upgrade: websocket\r\n"
                                "Connection: Upgrade\r\n"
                                "Sec-WebSocket-Accept: " + acceptKey + "\r\n\r\n").toLatin1());
        v_dataStream.clear();
        v_state = Opened;
        Q_EMIT handshakeSuccess();
        break;
    case Opened:
        do
        {
            frame = computeFrame(v_dataStream);
	        //qDebug() << "opcode=" << (int)(frame.opcode & 0xFF) << ", length=" << frame.dataIndex << "+" << frame.length;

            if(frame.opcode & 0x1 || frame.opcode & 0x2)
            {
                encodedData = v_dataStream.mid(frame.dataIndex, frame.length);
                maskCurrentIndex = 0;
                for(int i=0; i<encodedData.count(); i++){
                    decodedData.append(encodedData.at(i) ^ frame.maskKey[maskCurrentIndex % 4]);
                    maskCurrentIndex++;
                }

                Q_EMIT dataReceived(decodedData);
            }
            else if(frame.opcode & 0x8) // connection close
            {
                decodedData = QByteArray();
                decodedData.append(136);
                v_state = Closed;
                QIODevice::write(decodedData);
                this->close();
                Q_EMIT onClosing();
            }
        v_dataStream = v_dataStream.mid(frame.dataIndex + frame.length, -1);
        //frame.dataIndex + frame.length > v_dataStream.length()
        }while(!v_dataStream.isEmpty());
        break;    
    case Closed:
        break;
    }
}



QString WebSocket::computeAcceptKey(QString key)
{
    key += "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    QByteArray hash = QCryptographicHash::hash ( key.toUtf8(), QCryptographicHash::Sha1 );
    return hash.toBase64();
}

WebSocket::Frame WebSocket::computeFrame(QByteArray frame)
{
    Frame f;
    int maskIndex;
    int len;

    f.fin = (frame.at(0) & 0x80) >> 7;
    f.rsv1 = (frame.at(0) & 0x40) >> 6;
    f.rsv2 = (frame.at(0) & 0x20) >> 5;
    f.rsv3 = (frame.at(0) & 0x10) >> 4;
    f.opcode = (frame.at(0) & 0xF);
    f.mask = (frame.at(1) & 0x80) >> 7;
    len = (frame.at(1) & 0x7F);
    f.length = len;

    if(len <= 125){
        maskIndex = 2;
    }else if(len == 126){
        maskIndex = 4;
    }else if(len == 127){
        maskIndex = 10;
    }

    f.maskKey[0] = frame.at(maskIndex);
    f.maskKey[1] = frame.at(maskIndex+1);
    f.maskKey[2] = frame.at(maskIndex+2);
    f.maskKey[3] = frame.at(maskIndex+3);

    f.dataIndex = maskIndex + 4;

    return f;
}

